#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <math.h>
#include "common_omp.h"
#include <omp.h>

// We use vectors to store bins, so we do not need to deal
// with storage problems
#include <vector>
using namespace std;

//  serial program, with the binning idea

#define IDX(i,j,NB) ((i)*(NB)+(j))
#define OUTBOUND(i,j,NB) ((i) < 0 || (i) >= (NB) || (j) < 0 || (j) >= (NB))

inline int getBinID(particle_t & particle, int NB)
{
    return floor(particle.x / cutoff) * NB + floor(particle.y / cutoff); 
}

extern double size;

// note: bins is passed by reference, not by value.
void assign_bins(int NB, int n, particle_t * particles,  vector<vector<particle_t*> > & bins)
{
    // clear bins
    #pragma omp parallel for
    for (int k = 0; k < bins.size(); k ++)
    {
        bins[k].clear();
    }
    // is this parallel-safe?
    #pragma omp parallel for
    for (int k = 0; k < n; k ++)
    {
        bins[getBinID(particles[k], NB)].push_back(particles+k); 
    }
}

void apply_force_binned(int NB, vector<vector<particle_t*> > &bins, particle_t &particle, omp_lock_t &lock)
{
    int dstbinid = getBinID(particle, NB);
    int i = dstbinid / NB;
    int j = dstbinid % NB;
    int istart = max(i-1,0), iend = min(i+2,NB),
        jstart = max(j-1,0), jend = min(j+2,NB);
    //    #pragma omp parallel for
    for (int _i = istart; _i < iend; _i ++)
    { 
      //  #pragma omp parallel for
        for (int _j = jstart; _j < jend; _j ++)
        {
            vector<particle_t *> & neighborbin = bins[IDX(_i,_j,NB)];
            for (int src = 0; src < neighborbin.size(); src ++)
              { 
                apply_force(particle, *neighborbin[src]);
              } 
        }
    }

}


int main( int argc, char **argv )
{    
    if( find_option( argc, argv, "-h" ) >= 0 )
    {
        printf( "Options:\n" );
        printf( "-h to see this help\n" );
        printf( "-n <int> to set the number of particles\n" );
        printf( "-o <filename> to specify the output file name\n" );
        return 0;
    }
    
    int n = read_int( argc, argv, "-n", 1000 );
    omp_set_nested(1);
    char *savename = read_string( argc, argv, "-o", NULL );
    
    FILE *fsave = savename ? fopen( savename, "w" ) : NULL;
    particle_t *particles = (particle_t*) malloc( n * sizeof(particle_t) );
    set_size( n );
    omp_lock_t *locks = (omp_lock_t*) malloc( n * sizeof(omp_lock_t)); 
    init_particles( n, particles );
    for (int i=0; i<n; i++){   
      omp_init_lock(&locks[i]);   
    }
    // create bins: we will use this to store the index of objects
    // It is highly inefficient: we may well end up with a large vector
    // but whatever...
    int NB = floor( size / cutoff) + 1; // NB is the number of bins per edge
    vector<vector<particle_t*> > bins;
    bins.resize(NB*NB);
    assign_bins(NB, n, particles, bins);

    //  simulate a number of time steps
    //
    double simulation_time = read_timer( );
    // this for loop can't be parallelized because future steps depend on
    // previous ones!
    for( int step = 0; step < NSTEPS; step++ )
    {
        // compute forces
        // we need to use blocks
        // first, initialize
      // the following pragma seems to account for most of the speedup
      #pragma omp parallel for
        for (int k = 0; k < n; k ++)
        {
            particles[k].ax = particles[k].ay = 0;
            apply_force_binned(NB, bins, particles[k], locks[k]);
        }
        
        //
        //  move particles
        //
        // Yangqing's note: for serial code, we do not need to 
        // change this at all since it runs in linear time

// TODO: do we need a barrier here??
#pragma omp parallel for
        for( int i = 0; i < n; i++ )
        {
            move( particles[i] );
        }

        // update bins again
        assign_bins(NB, n, particles, bins);

        //
        //  save if necessary
        //
#pragma omp master
        if( fsave && (step%SAVEFREQ) == 0 )
            save( fsave, n, particles );
    }
    simulation_time = read_timer( ) - simulation_time;
    
    printf( "n = %d, simulation time = %g seconds\n", n, simulation_time );
    
    free( locks );
    free( particles );
    if( fsave )
        fclose( fsave );
    
    return 0;
}

